use anyhow::anyhow;
use polars::prelude::{
    DataFrame, IntoLazy, JsonFormat, JsonReader, LazyFrame, SerReader, UnionArgs, concat,
};
use serde::Serializer;
use serde::ser::SerializeMap;
use serde_json::Value;
use std::io::Write;
use std::num::NonZeroUsize;

pub fn convert_to_ndjson(
    mut buffer: Vec<u8>,
    fields: Vec<String>,
    items: Vec<Vec<Value>>,
) -> anyhow::Result<Vec<u8>> {
    let len = fields.len();
    // 在独立作用域中完成写入
    {
        let mut writer = std::io::BufWriter::new(&mut buffer);
        for row in &items {
            let mut ser = serde_json::Serializer::new(&mut writer);
            let mut map = ser.serialize_map(Some(len))?;

            for (idx, field) in fields.iter().enumerate() {
                let val = row.get(idx).unwrap_or(&Value::Null);
                map.serialize_entry(field, val)?;
            }
            map.end()?;
            writer.write_all(b"\n")?;
        }

        // 显式 flush 确保数据写入
        writer.flush()?;
    } // 这里 writer 被 drop，释放对 buffer 的借用

    Ok(buffer) // 现在可以安全返回所有权
}

pub fn build_pl_df(buffer: Vec<u8>) -> anyhow::Result<DataFrame> {
    let schema_len = NonZeroUsize::new(1000usize)
        .ok_or(anyhow!("invalid schema length"))?;
    let df = JsonReader::new(std::io::Cursor::new(buffer))
        .with_json_format(JsonFormat::JsonLines)
        .infer_schema_len(Some(schema_len))
        .finish()?;
    Ok(df)
}

pub fn concat_dfs(dfs: Vec<DataFrame>) -> anyhow::Result<LazyFrame> {
    let non_empty: Vec<_> = dfs
        .into_iter()
        .filter(|df| !df.is_empty())
        .map(|df| df.lazy())
        .collect();
    concat(non_empty, UnionArgs::default()).map_err(|e| anyhow!(e))
}

#[macro_export]
macro_rules! get_str_col {
    ($row:expr, $idx:expr) => {
        $row.0[$idx].str_value()
    };
}

#[macro_export]
macro_rules! get_u32_col {
    ($row:expr, $idx:expr) => {
        $row.0[$idx].try_extract::<u32>()?
    };
}

#[macro_export]
macro_rules! get_f64_col {
    ($row:expr, $idx:expr) => {
        $row.0[$idx].try_extract::<f64>()?
    };
}

#[macro_export]
macro_rules! get_col_idx {
    ($df:expr) => {{
        $df.schema()
            .iter_names()
            .enumerate()
            .map(|(i, f)| (f.as_str(), i))
            .collect::<std::collections::HashMap<&str, usize>>()
    }};
}
